/**
* This file is part of ArmarX.
*
* ArmarX is free software; you can redistribute it and/or modify
* it under the terms of the GNU Lesser General Public License as
* published by the Free Software Foundation; either version 2 of
* the License, or (at your option) any later version.
*
* ArmarX is distributed in the hope that it will be useful, but
* WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
* GNU Lesser General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*
* @package    ArmarX::
* @author     Mirko Waechter ( mirko.waechter at kit dot edu)
* @date       2014
* @copyright  http://www.gnu.org/licenses/gpl.txt
*             GNU General Public License
*/

#include "RigidBodyConverter.h"
#include <MMM/Motion/Legacy/LegacyMotion.h>
#include <MMMSimoxTools/MMMSimoxTools.h>
#include <VirtualRobot/RobotNodeSet.h>

using namespace MMM;
using namespace VirtualRobot;

RigidBodyConverter::RigidBodyConverter(const std::string& name) :
    ConverterVicon2MMM(name)
{
}




bool RigidBodyConverter::buildModel(MMM::ModelPtr model)
{
    mmmModel = MMM::SimoxTools::buildModel(model, false);
    if (!mmmModel)
        return false;

    if (modelName.empty())
        mmmModel->setName(mmmModel->getType());
    else
        mmmModel->setName(modelName);

    // store markes
    std::vector<SensorPtr> sensors = mmmModel->getSensors();
    for (size_t i = 0; i < sensors.size(); i++)
    {
        SensorPtr sensor = sensors[i];
        modelMarkerPositions[sensor->getName()] = sensor->getGlobalPose().block<3,1>(0,3);
    }
    return true;
}


AbstractMotionPtr RigidBodyConverter::initializeStepwiseConversion()
{
    if (!isInitialized())
    {
        MMM_WARNING << "not initialized..." << endl;
        return AbstractMotionPtr();
    }
    initialModelPose = mmmModel->getGlobalPose();
    LegacyMotionPtr res(new LegacyMotion(outputModel->getName()));
    return res;
}

bool RigidBodyConverter::convertMotionStep(AbstractMotionPtr currentOutput)
{
    if(!isInitialized())
        return false;
    if (!currentOutput)
    {
        return false;
    }

    if (currentOutput->getNumFrames()>=inputMarkerMotion->getNumFrames())
    {
        MMM_WARNING << "Could not proceed. End of input motion reached..." << endl;
        return false;
    }

    LegacyMotionPtr outputMotion = boost::dynamic_pointer_cast<LegacyMotion>(currentOutput);
    if (!outputMotion)
    {
        MMM_ERROR << "Expecting an outputMotion of type Motion..." << endl;
        return false;
    }

    unsigned int pos = currentOutput->getNumFrames();
    MMM::LegacyMarkerDataPtr frame = inputMarkerMotion->getFrame(pos);
    if (!frame)
    {
        MMM_ERROR << "Could not grab frame " << pos << " from input marker motion" << endl;
        return false;
    }

    // map input marker names onto output marker names
    std::map< std::string, Eigen::Vector3f > observedMarkers;
    for(std::map< std::string, Eigen::Vector3f >::const_iterator it = frame->data.begin(); it != frame->data.end(); it++)
    {
        if(!it->first.empty() && !markerMapping[it->first].empty())
            observedMarkers[markerMapping[it->first]] = it->second;
    }

    Eigen::Matrix4f rigidBodyPose = Registration3d(modelMarkerPositions,observedMarkers);

#ifdef CONVERTER_OUTPUT_MARKER_DEVIATION
    // Calculate and print error
    float totalD = 0.0f, maxD = 0.0f;
    int markerCount = 0;

    for(std::map< std::string, Eigen::Vector3f >::const_iterator it = modelMarkerPositions.begin(); it != modelMarkerPositions.end(); it++)
    {
        Eigen::Vector4f modelMarker;
        modelMarker.head(3) = it->second;
        modelMarker[3] = 1.0f;
        Eigen::Vector4f modelMarkerTransformed = rigidBodyPose * modelMarker;
        float distance = (modelMarkerTransformed.head(3) - observedMarkers[it->first]).norm();

        /* std::cout << "Marker name: " << it->first << std::endl;
        std::cout << "Physical marker:\n" << observedMarkers[it->first] << std::endl;
        std::cout << "Virtual marker:\n" << modelMarkerTransformed << std::endl;
        std::cout << "Distance: " << distance << std::endl; */

        ++markerCount;
        totalD += distance;
        if (distance > maxD)
            maxD = distance;
    }

    float avgD = markerCount ? (totalD / markerCount) : FLT_MAX;  // markerCount == 0 should never happen, but if it does, just print FLT_MAX
    std::cout << "Frame #" << frame->frame << " finished: max error = " << maxD << ", avg error = " << avgD << std::endl;
#endif

    MotionFramePtr md(new MotionFrame(0));
    md->timestep = frame->timestamp;
    md->setRootPose(rigidBodyPose);
    outputMotion->addMotionFrame(md);

    return true;
}

Eigen::Matrix4f RigidBodyConverter::Registration3d(const Eigen::MatrixXf &A, const Eigen::MatrixXf &B)
{

    // std::cout << "A:\n" << A<< std::endl;
    // std::cout << "B:\n" << B<< std::endl;
    assert(A.rows() == B.rows()); // Same number of points
    assert(A.cols() == B.cols()); // Same dimension
    //    assert(A.cols()==3);
    assert(A.cols() == 3 || A.cols() == 4); // Homogeneous coordinates?
    assert(A.rows() >= 3);


    Eigen::Vector3f centroid_A = A.block(0,0,A.rows(),3).colwise().mean();
    Eigen::Vector3f centroid_B = B.block(0,0,A.rows(),3).colwise().mean();
    // std::cout << "cA & cB: \n" << centroid_A << std::endl << std::endl << centroid_B << std::endl;

    Eigen::MatrixXf H = (A.block(0,0,A.rows(),3).rowwise() - centroid_A.transpose()).transpose() * (B.block(0,0,A.rows(),3).rowwise() - centroid_B.transpose());
    // std::cout << "H: " << H<< std::endl;

    Eigen::JacobiSVD<Eigen::MatrixXf> svd(H, Eigen::ComputeThinU | Eigen::ComputeThinV);
    Eigen::Matrix3f U = svd.matrixU();
    // std::cout << "U: " << U<< std::endl;
    Eigen::Matrix3f V = svd.matrixV();
    // std::cout << "V: " << V<< std::endl;
    Eigen::Matrix3f R = V*U.transpose();
    // std::cout << "R: " << R<< std::endl;

    if (R.determinant()<0) {
        // Reflection detected
//        abort();
        V.col(2) *= -1;
        R = V*U.transpose();
    }

    Eigen::Vector3f t = -R*centroid_A+centroid_B;
    // std::cout << "t: " << t << std::endl;
    Eigen::Matrix4f T = Eigen::Matrix4f::Identity();
    T.block(0,0,3,3) = R;
    T.block(0,3,3,1) = t;

    return T;
}


bool RigidBodyConverter::isInitialized()
{
    if(!inputMotion)
        return false;
    if(!outputModel)
        return false;
    if(!inputMarkerMotion)
        return false;
    if(!mmmModel)
        return false;
    if(inputMarkerMotion->getNumFrames()==0)
        return false;
    if(markerMapping.size() == 0)
        return false;
    return true;
}
